package com.yxc.elasticsearchdemo;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.tomcat.util.bcel.Const;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;

import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author: 杨旭晨
 * @time: 2022/6/23 14:49
 **/
@Slf4j
@SpringBootTest
public class UpdateById {

    @Autowired
    RestHighLevelClient client;

    @Test
    void updateById1() throws Exception {
        long startTime = 0L;
        long endTime = 0L;
        // 更新不同数据量的数据
        for (Integer dataSize : Constant.dataSizeList) {
            // 使用不用大小的快来更新
            for (Integer size : Constant.sizeList) {
                startTime = System.currentTimeMillis();
                // 构造查询条件
                SearchRequest searchRequest = new SearchRequest(Constant.esIndex)
                        .scroll(new Scroll(TimeValue.timeValueMinutes(1)))
                        .searchType(SearchType.DEFAULT)
                        .source(
                                new SearchSourceBuilder()
                                        .query(new MatchQueryBuilder("cName.keyword", "更新测试" + dataSize))
                                        .size(size)
                                        .fetchSource("", "")
                        );
                // 游标
                String scrollId = null;
                SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
                if (searchResponse.getHits().getHits().length > 0) {
                    // 命中，进行更新
                    SearchHit[] hits = searchResponse.getHits().getHits();
                    upateByIdBath(hits, "scroll更新" + dataSize + "块" + size);
                    scrollId = searchResponse.getScrollId();
                }
                // 当数据量大于分块大小时，继续循环更新
                if (searchResponse.getHits().getTotalHits().value > size) {
                    // 继续循环
                    while (StringUtils.isNotEmpty(scrollId)) {
                        SearchScrollRequest searchScrollRequest =
                                new SearchScrollRequest(scrollId).scroll(TimeValue.timeValueMinutes(1));
                        SearchResponse response = client.scroll(searchScrollRequest, RequestOptions.DEFAULT);
                        scrollId = response.getScrollId();
                        if (response.getHits().getHits().length > 0) {
                            // 命中，更新
                            SearchHit[] hits = response.getHits().getHits();
                            upateByIdBath(hits, "scroll更新" + dataSize + "块" + size);
                        } else {
                            scrollId = null;
                        }
                    }
                }
                endTime = System.currentTimeMillis();
                log.info("根据id更新，数据量{}，更新块大小{}，使用时间：{}", dataSize, size, endTime - startTime);
            }
        }
    }

    boolean upateByIdBath(SearchHit[] hits, String updateData) throws IOException {
        BulkRequest bulkRequest = new BulkRequest();
        for (SearchHit hit : hits) {
            Map<String, Object> doc = new HashMap<>();
            doc.put("cValue", updateData);
            UpdateRequest updateRequest = new UpdateRequest(Constant.esIndex, hit.getId()).doc(doc);
            bulkRequest.add(updateRequest);
        }
        BulkResponse bulk = client.bulk(bulkRequest, RequestOptions.DEFAULT);
        return true;
    }


}
